#!/usr/bin/env python3
"""
Convert trace XML exports to human-readable text format (ftrace-like)
Usage: ./xml_to_text.py <input.xml> [output.txt]
"""

import re
import argparse

def extract_timestamp(log):
    """Extract timestamp(us) from trace line"""
    m = re.search(r" (\d+)\.(\d+): ", log)
    if m is None:
        return 0
        
    match = m.group()
    stamp = int(match.strip().split('.')[0]) * int(1e6) + int(match.strip().rstrip(':').split('.')[-1])
    return stamp

def extract_thread_id(log):
    """Extract thread id from trace line"""
    m = re.search(" \(.+\)\s+\[\d", log)
    
    if m is None:
        return None
    
    return int(log[:m.span()[0]].split('-')[-1])

def extract_prcess_id(log):
    """Extract process id from trace line"""
    m = re.search(r"\(\s*\d+\) \[", log)
    
    if m is None:
        return 0
    
    match = m.group()
    if '-' in match:
        return 0
    return int(match.split(')')[0].lstrip('('))

def parse_sched_wakeup(log):
    """sched_wakeup comm=* pid=* prio=* target_cpu=*"""
    waker_tid = extract_thread_id(log)
    
    wakee_name = log[log.index(" comm=") + len(" comm="):log.index(" pid=")].strip()
    wakee_tid= int(log[log.index(" pid=") + len(" pid="):log.index(" prio=")])

    return waker_tid, wakee_name, wakee_tid

def parse_sched_switch(log):
    """
    sched_switch: prev_comm=* prev_pid=* prev_prio=* prev_state=* next_comm=* next_pid=* next_prio=*
    """
    prev_comm = log[log.index(" prev_comm=") + len(" prev_comm="):log.index(" prev_pid=")].strip()
    prev_tid = int(log[log.index(" prev_pid=") + len(" prev_pid="):log.index(" prev_prio=")])
    prev_state = log[log.index(" prev_state=") + len(" prev_state="):log.index(" ==>")].strip()
    
    next_comm = log[log.index(" next_comm=") + len(" next_comm="):log.index(" next_pid=")].strip()
    next_tid = int(log[log.index(" next_pid=") + len(" next_pid="):log.index(" next_prio=")])
    next_prio = int(log[log.index(" next_prio=") + len(" next_prio="):].split(' ')[0])
    
    return prev_comm, prev_tid, prev_state, next_comm, next_tid, next_prio

def parse_trace_records(file):
    start_time = None
    pid_map = {}
    thread_track = {}
    
    with open(file, 'r', encoding="gb18030", errors="ignore") as f:
        logs = f.readlines()

        for i in range(len(logs)):
            if "|I13" in logs[i]:
                logs[i] = logs[i][:logs[i].find("|I13")]
            if "|H:[" in logs[i] and "]#FFRT" in logs[i]:
                logs[i] = logs[i][:logs[i].find("|H:[") + 2] + logs[i][logs[i].find("]#FFRT") + 2:]
            log = logs[i]
        
            timestamp = extract_timestamp(log)
            if timestamp == 0:
                continue
            
            if start_time is None:
                start_time = timestamp
            
            pid = extract_prcess_id(log)
            if pid not in pid_map.keys():
                pid_map[pid] = []
            
            tid = extract_thread_id(log)
            if tid not in pid_map[pid]:
                pid_map[pid].append(tid)
                
            if "sched_wakeup:" in log:
                waker_tid, wakee_name, wakee_tid = parse_sched_wakeup(log)
                if wakee_tid not in thread_track.keys():
                    thread_track[wakee_tid] = []
                thread_track[wakee_tid].append([timestamp, "R"])
            elif "sched_switch" in log:
                prev_comm, prev_tid, prev_state, next_comm, next_tid, next_prio = parse_sched_switch(log)
                if prev_tid not in thread_track.keys():
                    thread_track[prev_tid] = []
                thread_track[prev_tid].append([timestamp, prev_state])
                if next_tid not in thread_track.keys():
                    thread_track[next_tid] = []
                thread_track[next_tid].append([timestamp, "Running"])
            elif "H:FFRT" in log:
                if tid not in thread_track.keys():
                    thread_track[tid] = []
                thread_track[tid].append([timestamp, log[log.find("H:FFRT"):]])
        f.close()

    return pid_map, thread_track, start_time

def generate_playback_records(file, pid):
    queue_info = {}
    task_info = {}
    
    pid_map, thread_track, start_time = parse_trace_records(file)

    if pid in pid_map.keys():
        threads = pid_map[pid]
        for thread in threads:
            track = thread_track[thread]
            
            task = None
            
            for e in track:
                timestamp = e[0]
                record  = e[1]
                
                if "H:FFRT" in record:
                    if "H:FFRTSubmit[" in record:
                        type = record.split('[')[1].split(']')[0]
                        
                        if type == "uv" or type == "nm":
                            gid = int(record.split('[')[2].split(']')[0])
                            qos = int(record.split('[')[3].split(']')[0])
                            task_info[gid] = {"type": type, "gid": gid, "qos": qos, "submit": timestamp, "execute": []}
                        elif type == "sq":
                            gid = int(record.split('[')[2].split(']')[0])
                            qos = int(record.split('[')[3].split(']')[0])
                            delay = int(record.split('[')[4].split(']')[0])
                            queue = int(record.split('[')[5].split(']')[0])
                            task_info[gid] = {"type": type, "gid": gid, "qos": qos, "delay": delay, "queue": queue, "submit": timestamp, "execute": []}
                            
                            if queue not in queue_info.keys():
                                queue_info[queue] = {"type":"sq"}
                        elif type == "cq":
                            gid = int(record.split('[')[2].split(']')[0])
                            qos = int(record.split('[')[3].split(']')[0])
                            delay = int(record.split('[')[4].split(']')[0])
                            queue = int(record.split('[')[5].split(']')[0])
                            concurrency = int(record.split('[')[6].split(']')[0])
                            task_info[gid] = {"type": type, "gid": gid, "qos": qos, "delay": delay, "queue": queue, "submit": timestamp, "execute": []}
                            
                            if queue not in queue_info.keys():
                                queue_info[queue] = {"type":"cq", "concurrency": concurrency}
                    if "H:FFRTExecute[" in record:
                        gid = int(record.split('[')[1].split(']')[0])
                        if gid in task_info.keys():
                            task_info[gid]["execute"].append([timestamp, "Running"])
                        task = gid
                    
                    if "H:FFRTFinish[" in record:
                        gid = int(record.split('[')[1].split(']')[0])
                        if gid in task_info.keys():
                            task_info[gid]["execute"].append([timestamp, "Finish"])
                        task = None
                    
                    if "H:FFRTCoWake[" in record:
                        gid = int(record.split('[')[1].split(']')[0])
                        qos = int(record.split('[')[2].split(']')[0])
                        if gid not in task_info.keys():
                            task_info[gid] = {"type": "nm", "gid": gid, "qos": qos, "submit": timestamp, "execute": []}
                        else:
                            task_info[gid]["execute"].append([timestamp, "Cowake"])
                else:
                    if task is not None:
                        if task in task_info.keys():
                            if record == "R" or record == "R+":
                                task_info[task]["execute"].append([timestamp, "Runnable"])
                            if record == "S" or record == "D":
                                task_info[task]["execute"].append([timestamp, "Sleeping"])
                            if record == "Running":
                                task_info[task]["execute"].append([timestamp, "Running"])
    
    for task in task_info.keys():
        commands = []
        states = sorted(task_info[task]["execute"], key=lambda x: x[0])
        for i in range(len(states) - 1):
            if states[i][1] == "Running" and states[i + 1][1] != "Running":
                commands.append([states[i + 1][0] - states[i][0], "Running"])
            if states[i][1] == "Sleeping" and states[i + 1][1] != "Sleeping":
                commands.append([states[i + 1][0] - states[i][0], "Sleeping"])
            if states[i][1] == "Finish" and states[i + 1][1] == "Cowake":
                commands.append([states[i + 1][0] - states[i][0], "Cowait"])
        task_info[task]["execute"] = commands
        task_info[task]["submit"] -= start_time

    with open(file + ".playback", "w") as file:
        for queue, info in queue_info.items():
            if info["type"] == "sq":
                file.write("queue,sq,%d,1,\n" % queue)
            else:
                file.write("queue,cq,%d,%d,\n" % (queue, info["concurrency"]))
        
        task_info_sorted = sorted(task_info.values(), key = lambda x: x["submit"])
        for info in task_info_sorted:
            if info["type"] == "uv" or info["type"] == "nm":
                if len(info["execute"]) > 0:
                    data = "task,%s,%d,%d,0,%d," % (info["type"], info["gid"], info["qos"], info["submit"])
                    for e in info["execute"]:
                        data += "%s,%d," % (e[1], e[0])
                    data += "\n"
                    file.write(data)

            if info["type"] == "sq" or info["type"] == "cq":
                if len(info["execute"]) > 0:
                    data = "task,%s,%d,%d,%d,%d," % (info["type"], info["gid"], info["qos"], info["queue"], info["submit"])
                    for e in info["execute"]:
                        data += "%s,%d," % (e[1], e[0])
                    data += "\n"
                    file.write(data)
            
        file.close()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="parse")
    parser.add_argument("--file", type=str, required=True, default=None, help="target trace file")
    parser.add_argument("--pid", type=int, required=True, default=None, help="specify pid to extrace")
    
    args = parser.parse_args()

    generate_playback_records(args.file, args.pid)
